#!/usr/bin/env python

import numpy as np
from PyQt4 import QtCore

floatTypes = (float,np.float64)
intTypes = (int,np.int32)

class numericArrayModel(QtCore.QAbstractTableModel):
    """
    Specific model for arrays of doubles and ints.
    """
    def __init__(self,parent=None):
        return QtCore.QAbstractTableModel.\
            __init__(self,parent)

    def __readArrayData__(self,index):
        if self.__theShape == 2:
            return self.__theArray[index.row(),index.column()]
        else:
            return self.__theArray[max(index.row(),index.column())]
    
    def __setArrayData__(self,index,value):
        if self.__theShape == 2:
            self.__theArray[index.row(),index.column()] = value
        else:
            self.__theArray[max(index.row(),index.column())] = value
    
    def __dataConversion__(self,value):
        if self.__theDataType in floatTypes:
            return value.toDouble()
        else:
            return value.toInt()

    def __writeOnCell__(self,index):
        if self.__theDataType in floatTypes:
            return QtCore.QVariant("%.6f" % self.__readArrayData__(index))
        else:
            return QtCore.QVariant("%d" % self.__readArrayData__(index))

    def setTheArray(self,array):
        self.__theArray = array
        self.__theShape = len(array.shape)
        self.__theDataType = None
        # Check if array dimension is not greater thant nxn
        if self.__theShape > 2:
            raise Exception("Array cannot be visualized\nbecause of its dimensionality.")
        # Check the array's element type
        if self.__theShape == 2:
            self.__theDataType = type(array[0,0])
        else:
            self.__theDataType = type(array[0])
        
    def rowCount(self,index):
        try:
            return self.__theArray.shape[0]
        except:
            return 1
    
    def columnCount(self,index):
        try:
            return self.__theArray.shape[1]
        except:
            return 1
    
    def headerData(self,section,orientation,role):
        if role == QtCore.Qt.DisplayRole:
            return QtCore.QVariant(section)
        return QtCore.QVariant()
        
    def data(self,index,role):
        if role == QtCore.Qt.TextAlignmentRole:
            return QtCore.QVariant(QtCore.Qt.AlignRight | QtCore.Qt.AlignVCenter)
        elif role in (QtCore.Qt.DisplayRole, QtCore.Qt.EditRole)\
            and index.isValid():
            return self.__writeOnCell__(index)
        return QtCore.QVariant()

    def flags(self,index):
        if not index.isValid():
            return QtCore.Qt.ItemIsEnabled
        flags = QtCore.QAbstractItemModel.flags(self,index)
        return flags | QtCore.Qt.ItemIsEditable

    def setData(self,index,value,role):
        if index.isValid() and role == QtCore.Qt.EditRole:
            value, conversionIsOk = self.__dataConversion__(value)
            if conversionIsOk: self.__setArrayData__(index,value)
            return True
        return False
